#!/usr/bin/env python

import cPickle
import optparse
import os
import sys

import minuit
import numpy
import matplotlib
from matplotlib import pyplot
import tables

import histogram
from charmicat import BadCommandLineError

import ypipi
 

def main():
    '''
    ypipi_mumupi0pi0_fits.py [opts] file.hpt [file.hpt [...]]

    Fit MM(pi0pi0) distributions for various cut combinations.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    opts, args = op.parse_args()

    if len(args) < 1:
        msg = 'Please supply at least one histogram file.'
        raise BadCommandLineError(msg, -1)

    dpath = dict([(os.path.basename(path), path) for path in args])

    fit_params = {}

    # Fits on the m_mumu ~ Y1S distributions
    path = dpath['hists_m_mumu_y1s_2sig.hpt']    
    h = histogram.load(path, names='mm_off_pi0pi0')
    
    # Y(1S) --> mu+ mu- pi0 pi0
    f = ypipi.fit.GaussPlusPolynomialFit(h, 4)
    f.fit(9.3, 9.7, signal_range=(9.4, 9.6))
    f.fit(9.3, 9.7, A=100, mu=9.46, sigma=0.05)
    f.plot(name="fit_y1s", title=r"$\Upsilon(1S)$ in $\Upsilon(5S)\to\Upsilon(1S)\pi^0\pi^0$")
    f.fix_signal()
    f.release_background()
    f.fit(9.3, 9.7)
    f.plot(name="fit_y1s_bg_released", title=r"$\Upsilon(1S)$ in $\Upsilon(5S)\to\Upsilon(1S)\pi^0\pi^0$")
    f.plot(name="fit_y1s_bg_subtracted", 
           title=r"$\Upsilon(1S)$ in $\Upsilon(5S)\to\Upsilon(1S)\pi^0\pi^0$",
           subtract_bg=True)

    fit_params['y1s'] = f.m.values

    # Y(2S) --> Y(1S)[mu+mu-] pi0 pi0
    f = ypipi.fit.GaussPlusPolynomialFit(h, 4)
    f.fit(9.9, 10.2, signal_range=(10, 10.1))
    f.fit(9.9, 10.2, A=100, mu=10.023, sigma=0.05)
    f.plot(name="fit_y2s_w_m_mumu_y1s_cut", title=r"$\Upsilon(2S)$ in $\Upsilon(2S)\to\Upsilon(1S)\pi^0\pi^0$")
    f.fix_signal()
    f.release_background()
    f.fit(9.9, 10.2)
    f.plot(name="fit_y2s_w_m_mumu_y1s_cut_bg_released", 
           title=r"$\Upsilon(2S)$ in $\Upsilon(2S)\to\Upsilon(1S)\pi^0\pi^0$")
    f.plot(name="fit_y2s_w_m_mumu_y1s_cut_bg_subtracted", 
           title=r"$\Upsilon(2S)$ in $\Upsilon(2S)\to\Upsilon(1S)\pi^0\pi^0$",
           subtract_bg=True)

    fit_params['y2s_y1s'] = f.m.values
    
    # Y(3S) --> Y(1S)[mu+mu-] pi0 pi0
    f = ypipi.fit.GaussPlusPolynomialFit(h, 4)
    f.fit(10.275, 10.5, signal_range=(10.3, 10.5))
    f.fit(10.275, 10.5, A=100, mu=10.35, sigma=0.05)
    f.plot(name="fit_y3s_w_m_mumu_y1s_cut", title=r"$\Upsilon(3S)$ in $\Upsilon(3S)\to\Upsilon(1S)\pi^0\pi^0$")
    f.fix_signal()
    f.release_background()
    f.fit(10.275, 10.55)
    f.plot(name="fit_y3s_w_m_mumu_y1s_cut_bg_released", 
           title=r"$\Upsilon(3S)$ in $\Upsilon(3S)\to\Upsilon(1S)\pi^0\pi^0$")
    f.plot(name="fit_y3s_w_m_mumu_y1s_cut_bg_subtracted", 
           title=r"$\Upsilon(3S)$ in $\Upsilon(3S)\to\Upsilon(1S)\pi^0\pi^0$",
           subtract_bg=True)

    fit_params['y3s_y1s'] = f.m.values

    # Fits on the m_mumu ~ Y2S distributions
    path = dpath['hists_m_mumu_y2s.hpt']    
    h = histogram.load(path, names='mm_off_pi0pi0')
    
    f = ypipi.fit.GaussPlusPolynomialFit(h, 4)
    f.fit(9.8, 10.2, signal_range=(10, 10.1))
    f.fit(9.8, 10.2, A=100, mu=10.023, sigma=0.05)
    f.plot(name="fit_y2s", title=r"$\Upsilon(2S)$ in $\Upsilon(5S)\to\Upsilon(2S)\pi^0\pi^0$")
    f.fix_signal()
    f.release_background()
    f.fit(9.8, 10.2)
    f.plot(name="fit_y2s_bg_released", title=r"$\Upsilon(2S)$ in $\Upsilon(5S)\to\Upsilon(2S)\pi^0\pi^0$")
    f.plot(name="fit_y2s_bg_subtracted", 
           title=r"$\Upsilon(2S)$ in $\Upsilon(5S)\to\Upsilon(2S)\pi^0\pi^0$",
           subtract_bg=True)

    fit_params['y2s'] = f.m.values

    # Fits on the m_mumu ~ Y3S distributions
    path = dpath['hists_m_mumu_y3s.hpt']    
    h = histogram.load(path, names='mm_off_pi0pi0')
    f = ypipi.fit.GaussPlusPolynomialFit(h, 4)
    f.fit(10.275, 10.55, signal_range=(10.35, 10.415))
    f.fit(10.275, 10.55, A=100, mu=10.35, sigma=0.05)
    f.plot(name="fit_y3s", title=r"$\Upsilon(3S)$ in $\Upsilon(5S)\to\Upsilon(3S)\pi^0\pi^0$")
    f.fix_signal()
    f.release_background()
    f.fit(10.275, 10.55)
    f.plot(name="fit_y3s_bg_released", 
           title=r"$\Upsilon(3S)$ in $\Upsilon(5S)\to\Upsilon(3S)\pi^0\pi^0$")
    f.plot(name="fit_y3s_bg_subtracted", 
           title=r"$\Upsilon(3S)$ in $\Upsilon(5S)\to\Upsilon(3S)\pi^0\pi^0$",
           subtract_bg=True)

    fit_params['y3s'] = f.m.values

    with open('fit_params.pkl', 'w') as fp:
        cPickle.dump(fit_params, fp)    


    return 0
    
    

if __name__ == "__main__":
    try:
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
